#!/usr/bin/env python
#
# htmldom.py
#
#  Copyright (c) 2005-2010  Yusuke Shinyama <yusuke at cs dot nyu dot edu>
#  
#  Permission is hereby granted, free of charge, to any person
#  obtaining a copy of this software and associated documentation
#  files (the "Software"), to deal in the Software without
#  restriction, including without limitation the rights to use,
#  copy, modify, merge, publish, distribute, sublicense, and/or
#  sell copies of the Software, and to permit persons to whom the
#  Software is furnished to do so, subject to the following
#  conditions:
#  
#  The above copyright notice and this permission notice shall be
#  included in all copies or substantial portions of the Software.
#  
#  THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY
#  KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE
#  WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
#  PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR
#  COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#  LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
#  OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
#  SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#

import sys
import urllib
from htmlparser3 import HTMLParser3, HTMLHandler
from style import Style
from utils import enc, first, last, attr2str
from utils import CDATA_TAGS, INLINE_IMMED_TAGS, FORM_FIELD_TAGS
from utils import NON_NESTED_TAGS, NON_DISPLAYED_TAGS, BREAK_LINE_TAGS

__all__ = [
  'HTMLElement', 'HTMLText', 'HTMLForm',
  'HTMLRootElement', 'HTMLDocumentBuilder',
  'parse_html', 'validate_dom'
  ]


##  HTMLElementBase
##
class HTMLElementBase(object):
  
  def __init__(self, root, tag, parent=None, attrs=None, active=True):
    "Creates an instance of the HTMLElementBase class."
    self.root = root
    self.tag = str(tag)  # not unicode object!
    self.parent = parent
    self.children = None
    if attrs:
      self._attrs = attrs.copy()
    else:
      self._attrs = {}
    self._active = active
    return
  
  def __getitem__(self, k):
    return self._attrs.get(k)

  def __setitem__(self, k, v):
    self._attrs[str(k)] = v
    return

  def get(self, *keys):
    return tuple( self._attrs.get(k) for k in keys )

  def dup(self):
    return HTMLElementBase(self.root, self.tag,
                           parent=self.parent, 
                           attrs=self._attrs,
                           active=self._active)

  def walk(self, cut=(lambda d,e: False), depth=0):
    if cut(depth, self): return
    depth += 1
    yield (depth, self)
    return
  
  def dump(self):
    raise NotImplementedError
  
  def add_child(self, e):
    raise TypeError('this element cannot have a child: %r' % self)
  
  def finish(self):
    self._active = False
    self._contexts = None
    return
  

##  HTMLText
##
class HTMLText(HTMLElementBase):

  def __init__(self, root, text, parent=None):
    HTMLElementBase.__init__(self, root, None, parent=parent)
    assert isinstance(text, basestring)
    self.text = text
    return

  def __repr__(self):
    return '<%r>' % self.text

  def __str__(self):
    return self.text

  def __iter__(self):
    return iter([])

  def dump(self, codec='ascii'):
    return enc(self.text, codec=codec)
  
  
##  HTMLElement
##
class HTMLElement(HTMLElementBase):
  
  """A object that represents an HTML element.
  
  Attributes:
  tag: the name of the element.
  root: the root element which is the whole HTML page.
  parent: the parent element.
  children: a list of child elements.
  attrs: dictionary which contains attributes.
  active: boolean flag which indicates whether this element is
          still under construction.
  """
  
  def __init__(self, root, tag, parent=None, children=None, attrs=None, active=True):
    "Creates an instance of the HTMLElement class."
    HTMLElementBase.__init__(self, root, tag, parent=parent, attrs=attrs, active=active)
    self.style = None
    if children is None:
      self.children = []
    else:
      self.children = children[:]
      for c in children:
        c.parent = self
    self._contexts = None
    return
  
  def __repr__(self):
    return '<%s%s>' % (self.tag, attr2str(self._attrs.iteritems()))

  def __iter__(self):
    return iter(self.children)

  def add_child(self, e):
    "Add a child element."
    self.children.append(e)
    e.parent = self
    return
  
  def dup(self):
    return HTMLElement(self.root, self.tag,
                       parent=self.parent, children=self.children, 
                       attrs=self._attrs, active=self._active)

  def walk(self, cut=(lambda d,e: False), depth=0):
    if cut(depth, self): return
    depth += 1
    yield (depth, self)
    if self.children:
      for c in self.children:
        for x in c.walk(cut=cut, depth=depth):
          yield x
    yield (-depth, self)
    return

  def dump(self, codec='ascii'):
    if self.tag == 'comment':
      yield '<!--'
      for c in self.children:
        for x in c.dump(codec=codec):
          yield x
      yield '-->'
    elif not self.children:
      attrs = self._attrs.iteritems()
      yield '<%s%s />' % (self.tag, attr2str(attrs, codec=codec))
    else:
      attrs = self._attrs.iteritems()
      yield '<%s%s>' % (self.tag, attr2str(attrs, codec=codec))
      for c in self.children:
        for x in c.dump(codec=codec):
          yield x
      yield u'</%s>' % self.tag
    return

  def finish(self):
    s = ''
    r = []
    for c in self.children:
      if isinstance(c, HTMLText):
        s += c.text
      else:
        if s:
          r.append(HTMLText(self.root, s, parent=self))
          s = ''
        r.append(c)
    if s:
      r.append(HTMLText(self.root, s, parent=self))
    self.children = r
    HTMLElementBase.finish(self)
    return
  
  def set_style(self, stylesheet):
    if not stylesheet: return
    attrs = self._attrs
    tags = [u'', self.tag]
    if 'class' in attrs:
      x = attrs['class']
      tags += [ t+'.'+x for t in tags ]
    if 'id' in attrs:
      x = attrs['id']
      tags += [ t+'#'+x for t in tags ]
    (style, contexts) = stylesheet.lookup(tags, self.parent._contexts)
    if 'style' in attrs:
      style = stylesheet.parse_style(attrs['style'], parent=style)
    self.style = style
    self._contexts = contexts
    return

  def get_text(self, ignore_tags=NON_DISPLAYED_TAGS, br_tags=BREAK_LINE_TAGS):
    for (d,e) in self.walk(cut=lambda _,e: e.tag in ignore_tags):
      if isinstance(e, HTMLText):
        yield e.text
      elif 0 < d:
        tag = e.tag
        if tag in br_tags:
          yield u'\n'
        if tag == 'img' and e['alt']:
          yield e['alt']
    return

  def get_links(self, normalize=False):
    """
    Generate the list of href links and its anchor text in the element.
    When normalize is set, the relative url is normalized based on the documentbase.
    """
    for (d,e) in self.walk():
      if 0 < d and e.tag == 'a' and e['href']:
        url = urllib.unquote(e['href'])
        if normalize:
          url = self.root.normalize_url(url)
        yield (e.get_text(), url)
    return


##  HTMLForm
##
class HTMLForm(object):
  
  def __init__(self, root, attrs):
    self.root = root
    self.attrs = attrs
    self.fields = []
    return
  
  def add_field(self, field):
    self.fields.append(field)
    field.form = self
    return


##  HTMLRootElement
##
class HTMLRootElement(HTMLElement):
  """A special element which contains the whole HTML document.
  """

  def __init__(self, charset=None, base_href=None):
    HTMLElement.__init__(self, self, 'html')
    self.charset = charset
    self.base_href = base_href
    self.forms = []
    self.head = HTMLElement(self, 'head')
    self.body = HTMLElement(self, 'body')
    return

  def dup(self):
    raise TypeError('HTMLRootElement cannot be duplicated.')

  def set_root_attrs(self, attrs):
    self.attrs = attrs
    return

  def set_root_charset(self, charset):
    self.charset = charset
    return

  def get_title(self):
    head = self.children[0]
    title = first( e for (d,e) in head.walk() if 0 < d and e.tag == 'title' )
    if not title: return None
    return ''.join(title.get_text())

  def add_header(self, e):
    self.head.add_child(e)
    if e.tag == 'base':
      if self.base_href:
        e['href'] = self.base_href
      else:
        self.base_href = e['href']
    return

  def add_child(self, e):
    self.body.add_child(e)
    return

  # normalize_url(url)
  def normalize_url(self, url):
    from urlparse import urljoin
    """Convert a relative URL to an absolute one."""
    return urljoin(self.base_href, url)

  def finish(self):
    if not first( e for e in self.head if e.tag == 'base' ):
      base = HTMLElement(self, 'base', attrs={ 'href': self.base_href })
      base.finish()
      self.head.add_child(base)
    self.head.finish()
    self.body.finish()
    HTMLElement.add_child(self, self.head)
    HTMLElement.add_child(self, self.body)
    HTMLElement.finish(self)
    return


##  HTMLDocumentBuilder
##
class HTMLDocumentBuilder(HTMLHandler):
  
  def __init__(self, base_href=None, stylesheet=None):
    self.root = HTMLRootElement(base_href=base_href)
    self.stylesheet = stylesheet
    self._form = None
    self._stack = [self.root]
    return

  def set_charset(self, charset):
    self.root.set_root_charset(charset)
    return

  def handle_data(self, s):
    assert self._stack
    self._stack[-1].add_child(HTMLText(self.root, s))
    return

  def do_unknown(self, tag, attrs):
    tag = str(tag)
    attrs = dict( (k.encode('ascii','ignore'),v) for (k,v) in attrs.iteritems() )
    e = HTMLElement(self.root, tag, attrs=attrs)
    if tag == 'link':
      self.root.add_header(e)
      if self.stylesheet and attrs.get('rel') == 'stylesheet' and 'href' in attrs:
        self.stylesheet.import_url(attrs.get('href'), attrs.get('media'))
    elif tag == 'base':
      self.root.add_header(e)
    elif tag == 'meta':
      self.root.add_header(e)
    else:
      self._stack[-1].add_child(e)
      e.set_style(self.stylesheet)
    return

  def start_unknown(self, tag, attrs):
    #print "start:", tag, attrs
    tag = str(tag)
    attrs = dict( (k.encode('ascii','ignore'),v) for (k,v) in attrs.iteritems() )
    if tag == 'html':
      self.root.set_root_attrs(attrs)
    elif tag == 'form':
      self._form = HTMLForm(self.root, attrs)
      self.root.forms.append(self._form)
    elif tag in NON_NESTED_TAGS:
      pass
    else:
      e = HTMLElement(self.root, tag, attrs=attrs)
      if tag == 'title':
        self.root.add_header(e)
      else:
        self._stack[-1].add_child(e)
        e.set_style(self.stylesheet)
      self._stack.append(e)
      if self._form and tag in FORM_FIELD_TAGS:
        self._form.add_field(e)
    return

  def end_unknown(self, tag):
    tag = str(tag)
    if tag == 'html':
      pass
    elif tag == 'form':
      self._form = None
    elif tag in NON_NESTED_TAGS:
      pass
    else:
      self._stack[-1].finish()
      if tag == 'style' and self.stylesheet:
        self.stylesheet.parse(''.join(self._stack[-1].get_text()))
        (style, _) = self.stylesheet.lookup(['',u'html',u'body'])
        self.root.style = style
      self._stack.pop()
    return

  def finish(self):
    assert len(self._stack) == 1 and self._stack[0] == self.root, self._stack
    self.root.finish()
    self._stack.pop()
    return



##  Utilities
##
def parse_html(x, base_href=None, charset=None, stylesheet=None):
  builder = HTMLDocumentBuilder(base_href=base_href, stylesheet=stylesheet)
  parser = HTMLParser3(builder, charset=charset)
  if isinstance(x, unicode):
    parser.feed_unicode(x)
    parser.close()
  elif isinstance(x, str):
    parser.feed_byte(x)
    parser.close()
  else:
    parser.feed_file(x)
    parser.close()
    x.close()
  return builder.root

def validate_dom(root, e0=None):
  "check if the structure is correct."
  # check if every node (except the root) has a parent.
  for (d,e) in (e0 or root).walk():
    if d < 0: continue
    if e != root and (not e.parent or (e not in iter(e.parent))):
      raise ValueError('Orphaned node: %r in %r' % (e, e.parent))
    # check if all my children belong to me.
    for c in e:
      if c.parent != e:
        raise ValueError('Invalid child: %r in %r' % (c, e))
  return


# main
if __name__ == '__main__':
  import getopt, urllib
  from style import ActiveStyleSheet
  def usage():
    print 'usage: htmldom.py [-d] [-b base_href] [-c charset_in] [-C codec_out] [url ...]'
    sys.exit(2)
  try:
    (opts, args) = getopt.getopt(sys.argv[1:], 'dc:C:b:')
  except getopt.GetoptError:
    usage()
  (charset_in, codec_out, base_href) = (None, 'utf-8', '')
  debug = 0
  for (k, v) in opts:
    if k == '-d': debug += 1
    elif k == '-c': charset_in = v
    elif k == '-C': codec_out = v
    elif k == '-b': base_href = v
  for url in args:
    fp = urllib.urlopen(url)
    stylesheet = ActiveStyleSheet()
    stylesheet.set_base(base_href or url)
    root = parse_html(fp, base_href=url, charset=charset_in,
                      stylesheet=stylesheet)
    fp.close()
    validate_dom(root)
    stylesheet.dump()
    if debug:
      for (d,e) in root.walk():
        if d < 0: continue
        sys.stdout.write('  '*d+repr(e)+'\n')
    else:
      for x in root.dump(codec=codec_out):
        sys.stdout.write(x)
